import torch
from transformers import AutoModelForImageClassification, AutoTokenizer, pipeline
from PIL import Image
import random

class CharacterSystem:
    def __init__(self, character_config):
        self.config = character_config
        self.dialogue_history = []
        self.emotional_state = "neutral"
        self.current_focus = None
        
    def update_state(self, visual_input, dialogue_context):
        """更新角色状态"""
        # 根据视觉输入和对话更新情感状态
        self.emotional_state = self._analyze_emotional_impact(visual_input, dialogue_context)
        # 更新当前关注点
        self.current_focus = self._determine_focus(visual_input, dialogue_context)
        
    def _analyze_emotional_impact(self, visual_input, dialogue_context):
        """分析情感影响"""
        # 基于角色特征分析情感反应
        emotional_responses = self.config["behavioral_patterns"]["emotional_responses"]
        for trigger, response in emotional_responses.items():
            if trigger in dialogue_context.lower():
                return response
        return "neutral"
        
    def _determine_focus(self, visual_input, dialogue_context):
        """确定当前关注点"""
        # 根据角色特征和知识库确定关注重点
        focus_points = []
        # 添加来自视觉输入的关注点
        if visual_input:
            focus_points.extend(visual_input.split(", "))
        # 添加来自对话的关注点
        if dialogue_context:
            for trait in self.config["core_traits"]["personality"]:
                if trait in dialogue_context:
                    focus_points.append(trait)
        return focus_points[0] if focus_points else None

    def generate_response_template(self, visual_input, query):
        """生成角色响应模板"""
        # 选择合适的响应模板
        templates = []
        
        # 基于性格特征的模板
        personality_templates = {
            "勇敢": [
                "作为一个格兰芬多，我必须说{}",
                "这让我想起了在霍格沃茨面对挑战时{}",
            ],
            "分析型": [
                "根据我的观察，{}",
                "从这些细节可以推断出{}",
            ]
        }
        
        # 根据当前角色特征选择模板
        for trait in self.config["core_traits"]["personality"]:
            if trait in personality_templates:
                templates.extend(personality_templates[trait])
                
        # 如果没有匹配的特征模板，使用通用模板
        if not templates:
            templates = [
                "关于这个问题，{}",
                "让我来回答这个问题，{}"
            ]
            
        # 随机选择一个模板
        template = random.choice(templates)
        
        return template

class MultimodalUnderstanding:
    def __init__(self):
        self.visual_memory = []
        self.context_history = []
        
    def analyze_visual_input(self, image_desc):
        """分析视觉输入"""
        # 处理图像描述
        visual_elements = image_desc.split(", ")
        
        # 提取关键视觉特征
        visual_features = {
            "objects": visual_elements,
            "scene_type": self._determine_scene_type(visual_elements),
            "emotional_tone": self._analyze_emotional_tone(visual_elements)
        }
        
        # 更新视觉记忆
        self.visual_memory.append(visual_features)
        return visual_features
        
    def _determine_scene_type(self, elements):
        """确定场景类型"""
        scene_keywords = {
            "indoor": ["room", "interior", "furniture"],
            "outdoor": ["sky", "tree", "street"],
            "abstract": ["pattern", "texture", "shape"]
        }
        
        for scene_type, keywords in scene_keywords.items():
            if any(keyword in " ".join(elements).lower() for keyword in keywords):
                return scene_type
        return "unknown"
        
    def _analyze_emotional_tone(self, elements):
        """分析情感基调"""
        emotion_keywords = {
            "positive": ["bright", "happy", "smile"],
            "negative": ["dark", "sad", "angry"],
            "neutral": ["plain", "simple", "regular"]
        }
        
        for tone, keywords in emotion_keywords.items():
            if any(keyword in " ".join(elements).lower() for keyword in keywords):
                return tone
        return "neutral"

class MMRoleLite:
    def __init__(self, config):
        self.config = config
        self.setup_models()
        self.character_systems = {}
        self.multimodal_understanding = MultimodalUnderstanding()
        self._initialize_character_systems()
        
    def _initialize_character_systems(self):
        """初始化角色系统"""
        for char_id, char_config in self.config.CHARACTERS.items():
            self.character_systems[char_id] = CharacterSystem(char_config)
        
    def setup_models(self):
        """初始化模型"""
        try:
            # 初始化视觉模型
            self.vision_model = pipeline(
                "image-classification",
                model=self.config.MODEL_CONFIG["vision_model"],
                device=self.config.MODEL_CONFIG["device"]
            )
            
            # 初始化文本生成模型
            self.text_generator = pipeline(
                "text-generation",
                model=self.config.MODEL_CONFIG["text_model"],
                device=self.config.MODEL_CONFIG["device"],
                trust_remote_code=True,
                framework="pt"
            )
            print("模型加载成功！")
        except Exception as e:
            print(f"模型加载失败: {str(e)}")
            raise
        
    def analyze_image(self, image_path):
        """图像分析"""
        try:
            image = Image.open(image_path)
            results = self.vision_model(image)
            visual_desc = results[0]['label'] if results else "无法识别的图像"
            # 使用多模态理解系统分析视觉输入
            visual_features = self.multimodal_understanding.analyze_visual_input(visual_desc)
            # 只返回描述文本，特征信息保存在实例中
            return visual_desc
        except Exception as e:
            print(f"图像分析错误: {str(e)}")
            return "图像处理出错"
    
    def generate_response(self, character_id, image_desc, user_input):
        """生成角色响应"""
        try:
            # 获取角色系统
            character_system = self.character_systems[character_id]
            character = self.config.CHARACTERS[character_id]
            
            # 更新角色状态
            character_system.update_state(image_desc, user_input)
            
            # 获取响应模板
            template = character_system.generate_response_template(image_desc, user_input)
            
            # 构建提示
            prompt = self._build_prompt(character, image_desc, user_input, template)
            
            # 生成回答
            response = self.text_generator(
                prompt,
                max_length=150,
                num_return_sequences=1,
                temperature=0.7,
                top_p=0.9,
                do_sample=True,
                repetition_penalty=1.2,
                pad_token_id=50256
            )
            
            generated_text = response[0]['generated_text']
            # 清理和验证生成的文本
            cleaned_text = self._clean_response(generated_text)
            
            # 验证文本是否合理
            if not self._is_valid_response(cleaned_text, character):
                return self._generate_fallback_response(character['name'], image_desc)
                
            return cleaned_text
            
        except Exception as e:
            print(f"响应生成错误: {str(e)}")
            return "生成回应时出错"

    def _build_prompt(self, character, image_desc, user_input, template):
        """构建提示词"""
        # 获取角色核心特征
        personality = "、".join(character["core_traits"]["personality"])
        values = "、".join(character["core_traits"]["values"])
        
        # 获取角色知识
        background = " ".join(character["knowledge_base"]["background"])
        
        # 构建提示词
        prompt = f"""作为{character['name']}，我的性格特点是{personality}，
我重视{values}。
{background}。

我看到一张图片，显示了{image_desc}。

{template.format(f'关于"{user_input}"这个问题')}"""
        
        return prompt

    def _is_valid_response(self, text, character):
        """验证生成的文本是否符合角色特征"""
        if len(text) < 20 or len(text) > 200:
            return False
        
        # 检查是否包含必要元素
        must_include = character["interaction_rules"]["must_include"]
        if not any(element in text for element in must_include):
            return False
            
        # 检查是否包含禁止元素
        must_avoid = character["interaction_rules"]["must_avoid"]
        if any(element in text for element in must_avoid):
            return False
            
        # 检查是否包含过多非中文字符
        chinese_chars = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
        if chinese_chars < len(text) * 0.5:
            return False
            
        return True
        
    def _generate_fallback_response(self, character_name, image_desc):
        """生成后备回答"""
        fallback_templates = {
            "哈利·波特": [
                f"这张图片让我想起了在霍格沃茨的一些经历。我看到了{image_desc}，这让我感到很有趣。",
                f"作为一个巫师，我对{image_desc}这样的景象并不陌生。这让我想起了一些魔法世界的事情。"
            ],
            "夏洛克·福尔摩斯": [
                f"根据我的观察，这张图片显示了{image_desc}。从这些细节中，我可以推断出一些有趣的事实。",
                f"让我仔细分析一下。图片中的{image_desc}展示了一些值得注意的特征。"
            ]
        }
        
        if character_name in fallback_templates:
            return random.choice(fallback_templates[character_name])
        return f"我看到了{image_desc}，这确实很有趣。"

    def _clean_response(self, text):
        """清理生成的文本"""
        try:
            # 移除所有提示词
            markers = [
                "请回答问题：",
                "让我来回答这个问题，",
                "关于这个问题，",
                "作为",
                "我的性格特点是",
                "我重视",
            ]
            
            # 清理文本
            for marker in markers:
                if marker in text:
                    text = text.split(marker)[-1]
            
            # 基本清理
            text = text.strip()
            text = text.replace('\n', ' ').replace('  ', ' ')
            
            # 确保以适当的标点符号结束
            if not text[-1] in ['。', '！', '？', '!', '?']:
                text += '。'
            
            return text
            
        except Exception as e:
            print(f"回答清理错误: {str(e)}")
            return "回答处理出错"

    def evaluate_response(self, response, character_id, image_desc):
        """评估响应质量"""
        try:
            character = self.config.CHARACTERS[character_id]
            scores = {
                "character_consistency": 0.0,
                "multimodal_understanding": 0.0,
                "interaction_quality": 0.0
            }
            
            if response == "生成回应时出错":
                return scores
                
            # 评估角色一致性
            scores["character_consistency"] = self._evaluate_character_consistency(
                response, character
            )
            
            # 评估多模态理解
            scores["multimodal_understanding"] = self._evaluate_multimodal(
                response, image_desc
            )
            
            # 评估交互质量
            scores["interaction_quality"] = self._evaluate_interaction_quality(
                response, character
            )
            
            return scores
        except Exception as e:
            print(f"评估错误: {str(e)}")
            return {
                "character_consistency": 0.0,
                "multimodal_understanding": 0.0,
                "interaction_quality": 0.0
            }

    def _evaluate_character_consistency(self, response, character):
        """评估角色一致性"""
        score = 3.0  # 降低基础分
        
        # 评估性格一致性
        personality_score = 0
        for trait in character["core_traits"]["personality"]:
            if trait in response:
                personality_score += 0.3  # 降低单项加分
        score += min(personality_score, 2.0)
        
        # 评估知识一致性
        knowledge_score = 0
        for bg in character["knowledge_base"]["background"]:
            if any(key in response for key in bg.split()):
                knowledge_score += 0.4
        score += min(knowledge_score, 2.5)
        
        # 评估行为连贯性
        behavior_score = 0
        for style in character["behavioral_patterns"]["speech_style"]:
            if style in response:
                behavior_score += 0.4
        score += min(behavior_score, 2.5)
        
        return min(score, 10.0)

    def _evaluate_multimodal(self, response, image_desc):
        """评估多模态理解"""
        score = 2.0  # 降低基础分
        
        # 评估视觉相关性
        visual_score = 0
        image_keywords = image_desc.lower().split()
        keywords_found = 0
        for keyword in image_keywords:
            if keyword in response.lower():
                keywords_found += 1
        # 根据找到的关键词比例计分
        visual_score = (keywords_found / len(image_keywords)) * 3.0
        score += visual_score
        
        # 评估上下文整合
        context_score = 0
        if len(response) > 50:  # 回答够长
            context_score += 1.0
        if any(keyword in response.lower() for keyword in image_keywords):  # 包含图像元素
            context_score += 1.0
        if "这张图片" in response or "我看到" in response:  # 明确引用图片
            context_score += 1.0
        score += min(context_score, 3.0)
        
        # 评估情感对齐
        emotion_keywords = {
            "positive": ["开心", "有趣", "喜欢", "吸引", "精彩"],
            "negative": ["担心", "害怕", "危险", "困扰", "麻烦"],
            "neutral": ["观察", "注意", "发现", "看起来", "似乎"]
        }
        emotion_score = 0
        emotion_matches = 0
        for keywords in emotion_keywords.values():
            for keyword in keywords:
                if keyword in response:
                    emotion_matches += 1
        emotion_score = min(emotion_matches * 0.5, 2.0)
        score += emotion_score
        
        return min(score, 10.0)

    def _evaluate_interaction_quality(self, response, character):
        """评估交互质量"""
        score = 2.0  # 降低基础分
        
        # 评估回应自然度
        naturalness_score = 0
        if 30 <= len(response) <= 100:  # 适中的长度
            naturalness_score += 1.0
        elif len(response) > 100:  # 较长但不过长
            naturalness_score += 0.5
            
        if response.count("。") >= 2:  # 完整的句子结构
            naturalness_score += 0.5
        if response.count("，") >= 1:  # 使用标点符号
            naturalness_score += 0.5
        if not any(char in response for char in ["!", "！", "?", "？"] * 2):  # 不过度使用感叹号和问号
            naturalness_score += 0.5
        score += min(naturalness_score, 3.0)
        
        # 评估参与度
        engagement_score = 0
        must_include = character["interaction_rules"]["must_include"]
        elements_found = 0
        for element in must_include:
            if element in response:
                elements_found += 1
        # 根据包含必要元素的比例计分
        engagement_score = (elements_found / len(must_include)) * 3.0
        score += engagement_score
        
        # 评估适当性
        appropriateness_score = 0
        must_avoid = character["interaction_rules"]["must_avoid"]
        if not any(avoid in response for avoid in must_avoid):
            appropriateness_score += 1.0
        if not any(word in response for word in ["呃", "额", "那个", "这个"]):  # 避免口语化
            appropriateness_score += 0.5
        if not response.count("我") > 3:  # 避免过度自我中心
            appropriateness_score += 0.5
        score += min(appropriateness_score, 2.0)
        
        return min(score, 10.0) 